import cv2 as cv
import numpy as np
import torch
from torchvision import transforms

from config import device, im_size
from data_gen import data_transforms

if __name__ == '__main__':
    checkpoint = 'BEST_checkpoint.tar'
    checkpoint = torch.load(checkpoint)
    model = checkpoint['model']
    model = model.to(device)
    model.eval()

    transformer = data_transforms['valid']

    filename = 'input1.jpg'
    bgr_img = cv.imread(filename)
    # bg_h, bg_w = bgr_img.shape[:2]
    bgr_img = cv.resize(bgr_img, (im_size, im_size))

    x_test = torch.zeros((1, 3, im_size, im_size), dtype=torch.float)
    img = transforms.ToPILImage()(bgr_img)
    img = transformer(img)
    x_test[0:, 0:3, :, :] = img

    # print(x_test.size())

    with torch.no_grad():
        y_pred, _ = model(x_test)

    _, y_pred = y_pred.topk(1, 1, True, True)
    y_pred = y_pred.cpu().numpy()[0][0]
    print('y_pred.shape: ' + str(y_pred.shape))

    out = (y_pred * 255).astype(np.uint8)
    cv.imwrite('out.png', out)
